import kagglehub
import tensorflow as tf
from tensorflow.keras.layers import TFSMLayer
import numpy as np
from PIL import Image
import os

# 下载模型
path = kagglehub.model_download("google/mobilenet-v2/tensorFlow2/100-224-classification")
print("Path to model files:", path)

# 修改为使用TFSMLayer加载模型
model = tf.keras.Sequential([
    TFSMLayer(path, call_endpoint='serving_default')
])
print("模型加载成功!")

# 示例分类函数
def get_imagenet_labels():
    # 下载ImageNet标签文件
    labels_url = "https://storage.googleapis.com/download.tensorflow.org/data/ImageNetLabels.txt"
    labels_path = tf.keras.utils.get_file("ImageNetLabels.txt", labels_url)
    with open(labels_path) as f:
        labels = f.read().splitlines()
    return labels

# 在classify_image函数中使用
def classify_image(image_path):
    # 加载并预处理图像
    img = Image.open(image_path).resize((224, 224))
    img_array = np.array(img) / 255.0
    img_array = np.expand_dims(img_array, axis=0)
    
    # 进行预测 (结果现在是字典)
    predictions = model.predict(img_array)
    
    # 获取输出张量 (根据模型实际输出调整键名)
    output_key = list(predictions.keys())[0]  # 获取第一个输出键
    output = predictions[output_key]
    
    # 获取前5个预测结果
    top5_indices = np.argsort(output[0])[-5:][::-1]
    
    imagenet_labels = get_imagenet_labels()
    print("预测结果(前5个):")
    for i, idx in enumerate(top5_indices):
        label = imagenet_labels[idx] if idx < len(imagenet_labels) else f"未知类别({idx})"
        print(f"{i+1}. {label}, 置信度: {output[0][idx]:.4f}")

# 示例使用
if __name__ == "__main__":
    # 替换为您的测试图像路径
    test_image = os.path.join(os.getcwd(), "tests", "data", "image.png")
    classify_image(test_image)